from __future__ import with_statement

import sys
import logging
import bisect
import optparse
import itertools
import matplotlib
matplotlib.use("PS")

import pylab
import matplotlib.backends.backend_ps
import numpy
import scipy.cluster.vq
import mpl_toolkits.mplot3d

import common
import refolding
import refolding_common

logger = logging.getLogger("refolding_plots")
logger.addHandler(logging.StreamHandler())
logger.setLevel(logging.INFO)

def draw_scatterplot_data(all_events, filename):
    gridspec = matplotlib.pyplot.GridSpec(3, 3)
    figure = matplotlib.pyplot.figure(1, figsize=(8, 8))

    time_cl_subplotspec = gridspec.new_subplotspec((0, 0), colspan=3)
    time_cl_ax = figure.add_subplot(time_cl_subplotspec)
    x = [event.t * 1e3 for event in all_events]
    y = [event.unfolded_cl * 1e9 for event in all_events]
    time_cl_ax.scatter(x, y, s=1)
    time_cl_ax.set_xlabel("Time [ms]")
    time_cl_ax.set_ylabel("Contour length [nm]")

    time_z_subplotspec = gridspec.new_subplotspec((1, 0), colspan=3)
    time_z_ax = figure.add_subplot(time_z_subplotspec)
    x = [event.t * 1e3 for event in all_events]
    y = [event.z * 1e9 for event in all_events]
    time_z_ax.scatter(x, y, s=1)
    time_z_ax.set_xlabel("Time [ms]")
    time_z_ax.set_ylabel("Z [nm]")

    cl_z_subplotspec = gridspec.new_subplotspec((2, 0), colspan=3)
    cl_z_ax = figure.add_subplot(cl_z_subplotspec)
    x = [event.unfolded_cl * 1e9 for event in all_events]
    y = [event.z * 1e9 for event in all_events]
    cl_z_ax.scatter(x, y, s=1)
    cl_z_ax.set_xlabel("Contour length [nm]")
    cl_z_ax.set_ylabel("Z [nm]")

    figure.text(0.0625, 0.91, "a)")
    figure.text(0.0625, 0.61, "b)")
    figure.text(0.0625, 0.31, "c)")

    canvas = matplotlib.backends.backend_ps.FigureCanvasPS(figure)
    canvas.print_figure(filename)

def draw_scatterplot_3d(
    loop_length, text, max_force, max_time, only_after_z, all_events, scatterplot_3d_filename):
    hovers = [event.to_hover() for event in all_events]
    x = refolding.compute_forces(refolding_common.kb_t_over_p, hovers)
    y = numpy.array([event.t for event in all_events], dtype=refolding_common.float_type)
    z = numpy.array(
        [float(event.folded) / (
                event.to_timed_hover().simplified_unfolded(loop_length) if only_after_z else event.to_timed_hover().unfolded()) for event in all_events],
        dtype=refolding_common.float_type)

    with open(scatterplot_3d_filename + ".txt", "wt") as f:
        for (xx, yy, zz) in zip(x, y, z):
            f.write("%r %r %r\n" % (xx * 1e12, yy * 1000, zz * 100))

    # definitions for the axes
    left, width = 0.1, 0.65
    bottom, height = 0.1, 0.65
    bottom_h = left_h = left + width + 0.02
    rect_scatter = [left, bottom, width, height]
    rect_histx = [left, bottom_h, width, 0.2]
    rect_histy = [left_h, bottom, 0.2, height]

    figure = matplotlib.pyplot.figure(1, figsize=(8, 8))

    ax = mpl_toolkits.mplot3d.Axes3D(figure)
    colours = refolding_common.gradient(y)
    # Looks like scatter and scatter3D do the same thing.
    ax.scatter3D(
        x * 1e12, y * 1000, z * 100, marker="o",
        color=colours, label=map(repr, x), s=3
    )
    if text:
        for (xx, yy, zz, event, c) in zip(x, y, z, all_events, colours):
            #iaca = ax.text(xx * 1e12, yy * 1000, zz * 100, event.name, zdir=(1, 0, 0))
            t = matplotlib.axes.Axes.text(ax, xx * 1e12, yy * 1000, event.name.replace("_break", ""), color=c)
            mpl_toolkits.mplot3d.art3d.text_2d_to_3d(t, zz * 100, zdir=(1, 0, 0))
    ax.set_xlabel("F [pN]")
    ax.set_ylabel("t [ms]")
    ax.set_zlabel("%")
    if max_force is None:
        max_force = max(x)
    else:
        common.check_ge(max_force, max(x))
    ax.set_xlim3d(0, max_force * 1e12)
    if max_time is None:
        max_time = max(y)
    else:
        common.check_ge(max_time, max(y))
    ax.set_ylim3d(0, max_time * 1000)
    ax.set_zlim3d(0, 100)

    canvas = matplotlib.backends.backend_ps.FigureCanvasPS(figure)
    canvas.print_figure(scatterplot_3d_filename)

def cluster_analysis(features, clusters):
    # Will use std_dev to compute the centroids in the original space.
    std_dev = numpy.std(features, axis=0)
    white_features = features / std_dev
    scipy.cluster.vq.whiten(white_features)
    code_book, distortion = scipy.cluster.vq.kmeans(white_features, clusters)
    logger.debug("distortion=%f code_book=%r", distortion, code_book * std_dev)
    code, dist = scipy.cluster.vq.vq(white_features, code_book)
    cluster_members = map(lambda (x, y): (x, list(y)), itertools.groupby(sorted(code)))
    common.check_equal(list(zip(*cluster_members)[0]), range(max(code) + 1))
    code_size_list = zip(code_book * std_dev, map(len, zip(*cluster_members)[1]))
    return code_size_list

def main(params):
    parser = optparse.OptionParser()
    refolding.add_reading_options(parser)
    parser.add_option("--scatterplot_3d_filename", help="eps filename") 
    parser.add_option("--loop_nm", help="loop length")
    parser.add_option("--text", choices=("True", "true", "False", "false"))
    parser.add_option("--max_force_pn", help="automatically determined if empty")
    parser.add_option("--max_time_ms", help="automatically determined if empty")
    parser.add_option("--only_after_z", choices=("True", "true", "False", "false"))
    parser.add_option("--scatterplot_data_filename", help="eps filename")
    parser.add_option("--clusters", help="number of clusters")

    options, args = parser.parse_args(params)
    if args:
        parser.print_help()
        return -1
    print(options)

    ignore_starred  = refolding.is_ignore_starred(options)
    all_events = refolding.read_all_events(options.root_dirs, ignore_starred)
    loop_length = refolding_common.float_type(options.loop_nm) / 1e9
    scatterplot_3d_filename = options.scatterplot_3d_filename
    text = {"false" : False, "true" : True}.get(options.text.lower())
    max_force = refolding_common.float_type(options.max_force_pn) / 1e12 if options.max_force_pn else None
    max_time = refolding_common.float_type(options.max_time_ms) / 1e3 if options.max_time_ms else None
    only_after_z = {"false" : False, "true" : True}.get(options.only_after_z.lower())
    scatterplot_data_filename = options.scatterplot_data_filename

    #logger.info([event for event in all_events if event.t < 0.01])
    clusters = int(options.clusters)
    features = numpy.array([[event.unfolded, event.z, event.t] for event in all_events])
    code_size_list = cluster_analysis(features, clusters)
    print(code_size_list)
    return

    #pylab.rc("font", family="serif")
    pylab.rc("font", size=10)
    #pylab.rc("axes", linewidth=0.5)
    #pylab.rc('legend', fontsize='small')

    if scatterplot_data_filename:
        draw_scatterplot_data(all_events, scatterplot_data_filename)

    if scatterplot_3d_filename:
        draw_scatterplot_3d(
            loop_length, text, max_force, max_time, only_after_z, all_events, scatterplot_3d_filename)


if __name__ == "__main__":
    main(sys.argv[1:])
